"""
本文件用于将使用segment.py分割好的目标掩膜移动到目标路径中去
"""

import os
import shutil
import argparse
from glob import glob


def move_masks_from_list(src_path, dst_dir, dry_run=False):
    """
    将分割好的掩膜文件从给定文件移动到目标目录，并重命名
    Args:
        src_path (str): 掩膜文件的完整路径
        dst_dir (str): 目标目录
        dry_run (bool): 是否只打印不实际移动
    """
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir, exist_ok=True)
        print(f"已创建目标目录: {dst_dir}")

    if not src_path:
        print("文件为空，未找到需要移动的文件")
        return

    # 例如：.../random_ds/seg_infer_R003_t2_tse_sag_p2_320_UNet2D_registered/prostate.nii.gz
    parent_dir = os.path.basename(os.path.dirname(src_path))  # seg_infer_xxx
    # 提取原文件名（如 R003_t2_tse_sag_p2_320_UNet2D_registered）
    if parent_dir.startswith("seg_infer_"):
        new_name = "prostate_" + parent_dir[len("seg_infer_") :] + ".nii.gz"
    else:
        new_name = "prostate_" + parent_dir + ".nii.gz"
    dst_path = os.path.join(dst_dir, new_name)
    if dry_run:
        print(f"[DRY RUN] Would move: {src_path} -> {dst_path}")
    else:
        shutil.move(src_path, dst_path)
        print(f"已移动: {src_path} -> {dst_path}")


if __name__ == "__main__":
    source_dir_path = "../segmentation/CHEN-REN-GENG_tra"
    source_file_name = "prostate.nii.gz"
    source_file_names = glob(
        os.path.join(source_dir_path, "**", source_file_name), recursive=True
    )
    print(f"Find total files: {len(source_file_names)}.")  # 12

    for file in source_file_names:
        # print(file)
        dst_dir = os.path.dirname(os.path.dirname(file))
        # print(dst_dir)
        # 想要把”/media/userdisk5/gqyang/final_result_t2map_/segmentations/cubic/random_ds/seg_infer_R003_t2_tse_sag_p2_320_UNet2D_registered/prostate.nii.gz“变为/media/userdisk5/gqyang/final_result_t2map_/segmentations/cubic/random_ds/prostate_infer_R003_t2_tse_sag_p2_320_UNet2D_registered.nii.gz作为目标文件路径
        move_masks_from_list(file, dst_dir, dry_run=False)
